
from pathlib import Path
import shutil
from classify.config import LabelManager
from classify.classifier import *
# ------------------------------
# 图片整理类
# ------------------------------
class ImageOrganizer:
    def __init__(self, src_folder, dst_folder, classifier:ImageClassifier, label_manager:LabelManager):
        self.src_folder = Path(os.path.expanduser(src_folder))
        self.dst_folder = Path(os.path.expanduser(dst_folder))
        self.dst_folder.mkdir(exist_ok=True)
        self.classifier = classifier
        self.label_manager = label_manager
        self.supported_suffixes = [".jpg", ".jpeg", ".png", ".bmp", ".gif"]

    def organize(self, recursive=True):
        pattern = "**/*" if recursive else "*"
        all_labels = self.label_manager.get_all_labels()

        # 一次性建立标签特征（如果已经有缓存会直接加载）
        self.classifier.build_label_features(all_labels)

        for file_path in self.src_folder.rglob(pattern):
            if not file_path.is_file():
                print(file_path, "is not a file")
                continue
            if file_path.suffix.lower() not in self.supported_suffixes:
                print(file_path, "is not a supported image")
                continue

            # 分类
            print(file_path, end="")
            try:
                category_label, score = self.classifier.classify(str(file_path))
            except Exception as e:
                print(f"⚠️ {file_path} 分类失败: {e}")
                bad_dir = self.dst_folder / "bad"
                bad_dir.mkdir(exist_ok=True)
                shutil.move(str(file_path), str(bad_dir / file_path.name))
                continue

            category_name = self.label_manager.get_name(category_label)

            # 创建目录并移动文件
            target_dir = self.dst_folder / category_name
            target_dir.mkdir(parents=True, exist_ok=True)
            shutil.move(str(file_path), str(target_dir / file_path.name))
            print(f" {file_path.name} → {category_name} ({score:.3f})")
